#pragma once

#include <cuda_runtime.h>

#include <utils/Print.h>

#include <string>
#include <type_traits>

namespace solar
{
    template <typename... Strings> auto stringConcatenate(Strings... strings) -> std::string
    {
        std::string result;
        (result.append(strings), ...);
        return result;
    }
} // namespace solar

#define CHECK_CUDA_ERRORS(val) solar::cuda::check((val), #val, __FILE__, __LINE__)

namespace solar::cuda
{
    inline void check(cudaError_t err, const char* const func, const char* const file,
                      const int line)
    {
        if (err != cudaSuccess)
        {
            solar::println(stringConcatenate("CUDA error at ", file, " : ", std::to_string(line)));
            solar::println(stringConcatenate("in function ", func, ": ", cudaGetErrorString(err)));
            throw "CUDA Error";
        }
    }

    template <typename T1, typename T2>
    inline auto cpu2gpu(T1*& d_out, T2* h_in, const size_t& size)
        -> std::enable_if_t<sizeof(T1) == sizeof(T2) && std::is_trivially_copy_assignable_v<T1> &&
                            std::is_trivially_copyable_v<T2>>
    {
        if (d_out == nullptr)
        {
            CHECK_CUDA_ERRORS(cudaMalloc(&d_out, sizeof(T1) * size));
        }
        CHECK_CUDA_ERRORS(cudaMemcpy(d_out, h_in, sizeof(T2) * size, cudaMemcpyHostToDevice));
    }

    template <typename T1, typename T2>
    inline auto cpu2gpuAsync(T1*& d_out, T2* h_in, const size_t& size,
                             const cudaStream_t& stream = nullptr)
        -> std::enable_if_t<sizeof(T1) == sizeof(T2) && std::is_trivially_copy_assignable_v<T1> &&
                            std::is_trivially_copyable_v<T2>>
    {
        if (d_out == nullptr)
        {
            CHECK_CUDA_ERRORS(cudaMallocAsync(&d_out, sizeof(T1) * size, stream));
        }
        CHECK_CUDA_ERRORS(
            cudaMemcpyAsync(d_out, h_in, sizeof(T2) * size, cudaMemcpyHostToDevice, stream));
    }

    // Copy from gpu to cpu
    // Note: you need to delete the host pointer manually if you provided a nullptr
    template <typename T1, typename T2>
    inline auto gpu2cpu(T1*& h_out, T2* d_in, const size_t& size)
        -> std::enable_if_t<sizeof(T1) == sizeof(T2) && std::is_trivially_copy_assignable_v<T1> &&
                            std::is_trivially_copyable_v<T2>>
    {
        if (h_out == nullptr)
        {
            h_out = new T1[size];
        }
        CHECK_CUDA_ERRORS(cudaMemcpy(h_out, d_in, sizeof(T2) * size, cudaMemcpyDeviceToHost));
    }

    // Copy from gpu to cpu
    // Note: you need to delete the host pointer manually if you provided a nullptr
    template <typename T1, typename T2>
    inline auto gpu2cpuAsync(T1*& h_out, T2* d_in, const size_t& size,
                             const cudaStream_t& stream = nullptr)
        -> std::enable_if_t<sizeof(T1) == sizeof(T2) && std::is_trivially_copy_assignable_v<T1> &&
                            std::is_trivially_copyable_v<T2>>
    {
        if (h_out == nullptr)
        {
            h_out = new T1[size];
        }
        CHECK_CUDA_ERRORS(cudaMemcpyAsync(h_out, d_in, sizeof(T2) * size, cudaMemcpyDeviceToHost, stream));
    }
} // namespace solar::cuda
